import java.time.Duration;
import java.time.Instant;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadLocalRandom;

public class ConcurrentMapTest {
    private final int threadNum;
    private String keyPrefix;
    private String valuePrefix;
    private ExecutorService threadPool;

    public ConcurrentMapTest(int threadNum) {
        this(threadNum, "key-mu-", "value-mu-");
    }

    public ConcurrentMapTest(int threadNum, String keyPrefix, String valuePrefix) {
        this.threadNum = threadNum;
        this.keyPrefix = keyPrefix;
        this.valuePrefix = valuePrefix;
        this.threadPool = Executors.newFixedThreadPool(threadNum);
        // warm up thread pool
        CountDownLatch latch = new CountDownLatch(threadNum);
        for (int i = 0; i < threadNum; i++) {
            threadPool.submit(latch::countDown);
        }
        try {
            latch.await();
        } catch (InterruptedException ex) {
            ex.printStackTrace();
        }
    }

    public Duration simulatePut(int opNum, ConcurrentMap<String, String> map,
                                boolean checkValue) throws InterruptedException {
        return simulatePut(opNum, map, checkValue, null);
    }

    public Duration simulatePut(int opNum, ConcurrentMap<String, String> map,
                                boolean checkValue, Integer collisionMod) throws InterruptedException {
        CountDownLatch latch = new CountDownLatch(threadNum);
        Instant startTime = Instant.now();
        for (int i = 0; i < threadNum; i++) {
            final int threadId = i;
            threadPool.submit(() -> {
                try {
                    int start = threadId * opNum;
                    int end = (threadId + 1) * opNum;
                    for (int j = start; j < end; j++) {
                        int suffix = j;
                        if (collisionMod != null)
                            suffix %= collisionMod;
                        String keyToPut = keyPrefix + suffix;
                        String valueToPut = valuePrefix + suffix;
                        map.put(keyToPut, valueToPut);
                        if (j > start && j % 7 == 0) {
                            int removeKey;
                            if (suffix == j)
                                removeKey = ThreadLocalRandom.current().nextInt(start, start + suffix);
                            else
                                removeKey = ThreadLocalRandom.current().nextInt(0, suffix + 1);
                            String keyToRemove = keyPrefix + removeKey;
                            String valueToRemove = map.remove(keyToRemove);
                            if (valueToRemove != null)
                                map.put(keyToRemove, valueToRemove);
                        }
                    }
                } catch (Exception ex) {
                    ex.printStackTrace();
                } finally {
                    latch.countDown();
                }
            });
        }
        latch.await();
        Instant endTime = Instant.now();
        Duration duration = Duration.between(startTime, endTime);
        if (checkValue) {
            System.out.println("Map size: " + map.size());
            for (int i = 0; i < threadNum * opNum; i++) {
                String key = keyPrefix + i;
                String value = map.get(key);
                if (!(valuePrefix + i).equals(value)) {
                    System.err.println("Assertion: " + key + ": " + value);
                }
            }
        }
        return duration;
    }

    public Duration simulateGet(int opNum, ConcurrentMap<String, String> map) throws InterruptedException {
        CountDownLatch latch = new CountDownLatch(threadNum);
        Instant startTime = Instant.now();
        final int startRange = 0;
        final int endRange = threadNum * opNum;
        for (int i = 0; i < threadNum; i++) {
            threadPool.submit(() -> {
                for (int j = 0; j < opNum; j++) {
                    int getKey = ThreadLocalRandom.current().nextInt(startRange, endRange);
                    String keyToGet = keyPrefix + getKey;
                    String valueToGet = map.get(keyToGet);
                    if (!(valuePrefix + getKey).equals(valueToGet)) {
                        System.err.println("Assertion: " + keyToGet + ": " + valueToGet);
                    }
                }
                latch.countDown();
            });
        }
        latch.await();
        Instant endTime = Instant.now();
        return Duration.between(startTime, endTime);
    }

    public void shutdown() {
        this.threadPool.shutdown();
    }
}
